from typing import List, Dict, Any, Optional
from datetime import datetime
from dataclasses import dataclass, field
import json
from utils.log_manager import get_system_logger

@dataclass
class QueryRecord:
    """单次查询记录"""
    timestamp: datetime
    user_query: str
    generated_sql: str
    result: str
    success: bool
    error_message: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        return {
            'timestamp': self.timestamp.isoformat(),
            'user_query': self.user_query,
            'generated_sql': self.generated_sql,
            'result': self.result,
            'success': self.success,
            'error_message': self.error_message,
            'metadata': self.metadata
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'QueryRecord':
        """从字典创建QueryRecord"""
        return cls(
            timestamp=datetime.fromisoformat(data['timestamp']),
            user_query=data['user_query'],
            generated_sql=data['generated_sql'],
            result=data['result'],
            success=data['success'],
            error_message=data.get('error_message'),
            metadata=data.get('metadata', {})
        )

@dataclass
class ContextMetadata:
    """上下文元数据"""
    session_id: str
    start_time: datetime
    last_activity: datetime
    total_queries: int = 0
    successful_queries: int = 0
    failed_queries: int = 0
    current_table_focus: Optional[str] = None  # 当前关注的表
    user_preferences: Dict[str, Any] = field(default_factory=dict)
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        return {
            'session_id': self.session_id,
            'start_time': self.start_time.isoformat(),
            'last_activity': self.last_activity.isoformat(),
            'total_queries': self.total_queries,
            'successful_queries': self.successful_queries,
            'failed_queries': self.failed_queries,
            'current_table_focus': self.current_table_focus,
            'user_preferences': self.user_preferences
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'ContextMetadata':
        """从字典创建ContextMetadata"""
        return cls(
            session_id=data['session_id'],
            start_time=datetime.fromisoformat(data['start_time']),
            last_activity=datetime.fromisoformat(data['last_activity']),
            total_queries=data.get('total_queries', 0),
            successful_queries=data.get('successful_queries', 0),
            failed_queries=data.get('failed_queries', 0),
            current_table_focus=data.get('current_table_focus'),
            user_preferences=data.get('user_preferences', {})
        )

class SessionContext:
    """会话上下文管理器"""
    
    def __init__(self, session_id: str, max_history: int = 20):
        """初始化会话上下文
        
        Args:
            session_id: 会话ID
            max_history: 最大历史记录数量
        """
        self.logger = get_system_logger('session_context')
        self.max_history = max_history
        
        # 初始化元数据
        self.metadata = ContextMetadata(
            session_id=session_id,
            start_time=datetime.now(),
            last_activity=datetime.now()
        )
        
        # 查询历史记录
        self.query_history: List[QueryRecord] = []
        
        self.logger.info("会话上下文初始化", session_id=session_id, max_history=max_history)
    
    def add_query_record(self, user_query: str, generated_sql: str, result: str, 
                        success: bool, error_message: Optional[str] = None,
                        metadata: Optional[Dict[str, Any]] = None) -> None:
        """添加查询记录"""
        record = QueryRecord(
            timestamp=datetime.now(),
            user_query=user_query,
            generated_sql=generated_sql,
            result=result,
            success=success,
            error_message=error_message,
            metadata=metadata or {}
        )
        
        self.query_history.append(record)
        
        # 更新元数据
        self.metadata.last_activity = datetime.now()
        self.metadata.total_queries += 1
        if success:
            self.metadata.successful_queries += 1
        else:
            self.metadata.failed_queries += 1
        
        # 维护历史记录数量限制
        if len(self.query_history) > self.max_history:
            removed_record = self.query_history.pop(0)
            self.logger.debug("移除旧查询记录", removed_query=removed_record.user_query[:50])
        
        self.logger.info("添加查询记录", 
                        user_query=user_query[:100],
                        success=success,
                        total_queries=self.metadata.total_queries)
    
    def get_recent_queries(self, count: int = 5) -> List[QueryRecord]:
        """获取最近的查询记录"""
        return self.query_history[-count:] if self.query_history else []
    
    def get_successful_queries(self, count: int = 5) -> List[QueryRecord]:
        """获取最近的成功查询记录"""
        successful = [record for record in self.query_history if record.success]
        return successful[-count:] if successful else []
    
    def get_context_summary(self) -> str:
        """获取上下文摘要，用于传递给LLM"""
        if not self.query_history:
            return "这是一个新的会话，没有历史查询记录。"
        
        recent_queries = self.get_recent_queries(3)
        context_parts = []
        
        context_parts.append(f"会话信息：已进行{self.metadata.total_queries}次查询，成功{self.metadata.successful_queries}次")
        
        if self.metadata.current_table_focus:
            context_parts.append(f"当前关注表：{self.metadata.current_table_focus}")
        
        if recent_queries:
            context_parts.append("最近的查询历史：")
            for i, record in enumerate(recent_queries, 1):
                status = "成功" if record.success else "失败"
                context_parts.append(f"{i}. 用户问题：{record.user_query}")
                context_parts.append(f"   生成SQL：{record.generated_sql}")
                context_parts.append(f"   查询结果：{record.result[:100]}..." if len(record.result) > 100 else f"   查询结果：{record.result}")
                context_parts.append(f"   状态：{status}")
        
        return "\n".join(context_parts)
    
    def update_table_focus(self, table_name: str) -> None:
        """更新当前关注的表"""
        self.metadata.current_table_focus = table_name
        self.logger.info("更新表关注", table_name=table_name)
    
    def extract_table_from_sql(self, sql: str) -> Optional[str]:
        """从SQL中提取表名"""
        try:
            sql_upper = sql.upper()
            if 'FROM' in sql_upper:
                from_index = sql_upper.find('FROM')
                after_from = sql[from_index + 4:].strip()
                # 简单提取第一个单词作为表名
                table_name = after_from.split()[0].strip('`"\'')
                return table_name
        except Exception as e:
            self.logger.warning("提取表名失败", sql=sql, error=str(e))
        return None
    
    def analyze_query_intent(self, user_query: str) -> Dict[str, Any]:
        """分析用户查询意图"""
        intent = {
            'is_follow_up': False,
            'references_previous': False,
            'contains_pronouns': False,
            'query_type': 'unknown'
        }
        
        query_lower = user_query.lower()
        
        # 检测跟进问题的关键词
        follow_up_keywords = ['再', '还', '另外', '接着', '然后', '继续', '也', '同样']
        intent['is_follow_up'] = any(keyword in query_lower for keyword in follow_up_keywords)
        
        # 检测代词引用
        pronouns = ['它', '这个', '那个', '这些', '那些', '上面', '前面', '刚才']
        intent['contains_pronouns'] = any(pronoun in query_lower for pronoun in pronouns)
        
        # 检测是否引用了之前的结果
        reference_keywords = ['结果', '数据', '记录', '刚才', '上次', '之前']
        intent['references_previous'] = any(keyword in query_lower for keyword in reference_keywords)
        
        # 简单的查询类型分类
        if any(word in query_lower for word in ['多少', '数量', 'count']):
            intent['query_type'] = 'count'
        elif any(word in query_lower for word in ['最大', '最小', '平均', 'max', 'min', 'avg']):
            intent['query_type'] = 'aggregation'
        elif any(word in query_lower for word in ['查找', '显示', '列出', 'select']):
            intent['query_type'] = 'select'
        elif any(word in query_lower for word in ['更新', '修改', 'update']):
            intent['query_type'] = 'update'
        
        return intent
    
    def clear_history(self) -> None:
        """清空查询历史"""
        self.query_history.clear()
        self.metadata.total_queries = 0
        self.metadata.successful_queries = 0
        self.metadata.failed_queries = 0
        self.metadata.current_table_focus = None
        self.logger.info("清空会话历史")
    
    def save_to_file(self, file_path: str) -> None:
        """保存会话到文件"""
        try:
            session_data = {
                'metadata': self.metadata.to_dict(),
                'query_history': [record.to_dict() for record in self.query_history],
                'max_history': self.max_history
            }
            
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(session_data, f, ensure_ascii=False, indent=2)
            
            self.logger.info("会话保存成功", file_path=file_path)
        except Exception as e:
            self.logger.error("会话保存失败", file_path=file_path, error=str(e))
    
    def load_from_file(self, file_path: str) -> None:
        """从文件加载会话"""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                session_data = json.load(f)
            
            self.metadata = ContextMetadata.from_dict(session_data['metadata'])
            self.query_history = [QueryRecord.from_dict(record) for record in session_data['query_history']]
            self.max_history = session_data.get('max_history', 20)
            
            self.logger.info("会话加载成功", 
                           file_path=file_path,
                           total_queries=self.metadata.total_queries)
        except Exception as e:
            self.logger.error("会话加载失败", file_path=file_path, error=str(e))
    
    def get_statistics(self) -> Dict[str, Any]:
        """获取会话统计信息"""
        return {
            'session_id': self.metadata.session_id,
            'duration': (self.metadata.last_activity - self.metadata.start_time).total_seconds(),
            'total_queries': self.metadata.total_queries,
            'successful_queries': self.metadata.successful_queries,
            'failed_queries': self.metadata.failed_queries,
            'success_rate': self.metadata.successful_queries / max(self.metadata.total_queries, 1),
            'current_table_focus': self.metadata.current_table_focus,
            'history_length': len(self.query_history)
        }